package com.chatbi.service.vectorstore;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.AllMiniLmL6V2QuantizedEmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import dev.langchain4j.data.document.Metadata;

import java.util.List;

@Component
public class ChromaVectorStore implements VectorStore {

    private final ChromaEmbeddingStore embeddingStore;
    private final EmbeddingModel embeddingModel;

    public ChromaVectorStore(@Value("${chroma.url:http://localhost:8000}") String chromaUrl) {
        this.embeddingStore = ChromaEmbeddingStore.builder()
                .baseUrl(chromaUrl)
                .collectionName("chatbi_vectors")
                .build();
        
        this.embeddingModel = new AllMiniLmL6V2QuantizedEmbeddingModel();
    }

    @Override
    public void store(String type, String content) {
        // Create embedding
        Embedding embedding = embeddingModel.embed(content).content();
        
        // Create metadata
        Metadata metadata = Metadata.from("type", type);
        
        // Create text segment
        TextSegment textSegment = TextSegment.from(content, metadata);
        
        // Store in Chroma
        embeddingStore.add(embedding, textSegment);
    }

    @Override
    public List<EmbeddingMatch<TextSegment>> loadEmbeddings(String type, String content, int batchSize, double minScore) {
        TextSegment textSegment = TextSegment.from(content, Metadata.from("type", type));
        Embedding queryEmbedding = embeddingModel.embed(textSegment).content();
        return embeddingStore.findRelevant(queryEmbedding, batchSize, minScore);
    }
}
